// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#include <aries/parquet/nebula/schema_internal.h>

#include <nebula/types/type.h>

using ArrowType = ::nebula::DataType;
using ArrowTypeId = ::nebula::Type;
using ParquetType = parquet::Type;

namespace parquet::arrow {
    using turbo::Result;
    using ::turbo::Status;
    using ::turbo::checked_cast;

    Result<std::shared_ptr<ArrowType> > MakeArrowDecimal(const LogicalType &logical_type) {
        const auto &decimal = checked_cast<const DecimalLogicalType &>(logical_type);
        if (decimal.precision() <= ::nebula::Decimal128Type::kMaxPrecision) {
            return ::nebula::Decimal128Type::create(decimal.precision(), decimal.scale());
        }
        return ::nebula::Decimal256Type::create(decimal.precision(), decimal.scale());
    }

    Result<std::shared_ptr<ArrowType> > MakeArrowInt(const LogicalType &logical_type) {
        const auto &integer = checked_cast<const IntLogicalType &>(logical_type);
        switch (integer.bit_width()) {
            case 8:
                return integer.is_signed() ? ::nebula::int8() : ::nebula::uint8();
            case 16:
                return integer.is_signed() ? ::nebula::int16() : ::nebula::uint16();
            case 32:
                return integer.is_signed() ? ::nebula::int32() : ::nebula::uint32();
            default:
                return turbo::invalid_argument_error(logical_type.ToString(),
                                                     " cannot annotate physical type Int32");
        }
    }

    Result<std::shared_ptr<ArrowType> > MakeArrowInt64(const LogicalType &logical_type) {
        const auto &integer = checked_cast<const IntLogicalType &>(logical_type);
        switch (integer.bit_width()) {
            case 64:
                return integer.is_signed() ? ::nebula::int64() : ::nebula::uint64();
            default:
                return turbo::invalid_argument_error(logical_type.ToString(),
                                                     " cannot annotate physical type Int64");
        }
    }

    Result<std::shared_ptr<ArrowType> > MakeArrowTime32(const LogicalType &logical_type) {
        const auto &time = checked_cast<const TimeLogicalType &>(logical_type);
        switch (time.time_unit()) {
            case LogicalType::TimeUnit::MILLIS:
                return ::nebula::time32(::nebula::TimeUnit::MILLI);
            default:
                return turbo::invalid_argument_error(logical_type.ToString(),
                                                     " cannot annotate physical type Time32");
        }
    }

    Result<std::shared_ptr<ArrowType> > MakeArrowTime64(const LogicalType &logical_type) {
        const auto &time = checked_cast<const TimeLogicalType &>(logical_type);
        switch (time.time_unit()) {
            case LogicalType::TimeUnit::MICROS:
                return ::nebula::time64(::nebula::TimeUnit::MICRO);
            case LogicalType::TimeUnit::NANOS:
                return ::nebula::time64(::nebula::TimeUnit::NANO);
            default:
                return turbo::invalid_argument_error(logical_type.ToString(),
                                                     " cannot annotate physical type Time64");
        }
    }

    Result<std::shared_ptr<ArrowType> > MakeArrowTimestamp(const LogicalType &logical_type) {
        const auto &timestamp = checked_cast<const TimestampLogicalType &>(logical_type);
        const bool utc_normalized = timestamp.is_adjusted_to_utc();
        static const char *utc_timezone = "UTC";
        switch (timestamp.time_unit()) {
            case LogicalType::TimeUnit::MILLIS:
                return (utc_normalized
                            ? ::nebula::timestamp(::nebula::TimeUnit::MILLI, utc_timezone)
                            : ::nebula::timestamp(::nebula::TimeUnit::MILLI));
            case LogicalType::TimeUnit::MICROS:
                return (utc_normalized
                            ? ::nebula::timestamp(::nebula::TimeUnit::MICRO, utc_timezone)
                            : ::nebula::timestamp(::nebula::TimeUnit::MICRO));
            case LogicalType::TimeUnit::NANOS:
                return (utc_normalized
                            ? ::nebula::timestamp(::nebula::TimeUnit::NANO, utc_timezone)
                            : ::nebula::timestamp(::nebula::TimeUnit::NANO));
            default:
                return turbo::invalid_argument_error("Unrecognized time unit in timestamp logical_type: ",
                                                     logical_type.ToString());
        }
    }

    Result<std::shared_ptr<ArrowType> > FromByteArray(const LogicalType &logical_type) {
        switch (logical_type.type()) {
            case LogicalType::Type::STRING:
                return ::nebula::utf8();
            case LogicalType::Type::DECIMAL:
                return MakeArrowDecimal(logical_type);
            case LogicalType::Type::NONE:
            case LogicalType::Type::ENUM:
            case LogicalType::Type::JSON:
            case LogicalType::Type::BSON:
                return ::nebula::binary();
            default:
                return turbo::unimplemented_error("Unhandled logical logical_type ",
                                                  logical_type.ToString(), " for binary array");
        }
    }

    Result<std::shared_ptr<ArrowType> > FromFLBA(const LogicalType &logical_type,
                                                 int32_t physical_length) {
        switch (logical_type.type()) {
            case LogicalType::Type::DECIMAL:
                return MakeArrowDecimal(logical_type);
            case LogicalType::Type::FLOAT16:
                return ::nebula::float16();
            case LogicalType::Type::NONE:
            case LogicalType::Type::INTERVAL:
            case LogicalType::Type::UUID:
                return ::nebula::fixed_size_binary(physical_length);
            default:
                return turbo::unimplemented_error("Unhandled logical logical_type ",
                                                  logical_type.ToString(),
                                                  " for fixed-length binary array");
        }
    }

    turbo::Result<std::shared_ptr<ArrowType> > FromInt32(const LogicalType &logical_type) {
        switch (logical_type.type()) {
            case LogicalType::Type::INT:
                return MakeArrowInt(logical_type);
            case LogicalType::Type::DATE:
                return ::nebula::date32();
            case LogicalType::Type::TIME:
                return MakeArrowTime32(logical_type);
            case LogicalType::Type::DECIMAL:
                return MakeArrowDecimal(logical_type);
            case LogicalType::Type::NONE:
                return ::nebula::int32();
            default:
                return turbo::unimplemented_error("Unhandled logical type ", logical_type.ToString(),
                                                  " for INT32");
        }
    }

    Result<std::shared_ptr<ArrowType> > FromInt64(const LogicalType &logical_type) {
        switch (logical_type.type()) {
            case LogicalType::Type::INT:
                return MakeArrowInt64(logical_type);
            case LogicalType::Type::DECIMAL:
                return MakeArrowDecimal(logical_type);
            case LogicalType::Type::TIMESTAMP:
                return MakeArrowTimestamp(logical_type);
            case LogicalType::Type::TIME:
                return MakeArrowTime64(logical_type);
            case LogicalType::Type::NONE:
                return ::nebula::int64();
            default:
                return turbo::unimplemented_error("Unhandled logical type ", logical_type.ToString(),
                                                  " for INT64");
        }
    }

    Result<std::shared_ptr<ArrowType> > GetArrowType(
        Type::type physical_type, const LogicalType &logical_type, int type_length,
        const ::nebula::TimeUnit::type int96_arrow_time_unit) {
        if (logical_type.is_invalid() || logical_type.is_null()) {
            return ::nebula::null();
        }

        switch (physical_type) {
            case ParquetType::BOOLEAN:
                return ::nebula::boolean();
            case ParquetType::INT32:
                return FromInt32(logical_type);
            case ParquetType::INT64:
                return FromInt64(logical_type);
            case ParquetType::INT96:
                return ::nebula::timestamp(int96_arrow_time_unit);
            case ParquetType::FLOAT:
                return ::nebula::float32();
            case ParquetType::DOUBLE:
                return ::nebula::float64();
            case ParquetType::BYTE_ARRAY:
                return FromByteArray(logical_type);
            case ParquetType::FIXED_LEN_BYTE_ARRAY:
                return FromFLBA(logical_type, type_length);
            default: {
                // PARQUET-1565: This can occur if the file is corrupt
                return turbo::io_error("Invalid physical column type: ",
                                       TypeToString(physical_type));
            }
        }
    }

    Result<std::shared_ptr<ArrowType> > GetArrowType(
        const schema::PrimitiveNode &primitive,
        const ::nebula::TimeUnit::type int96_arrow_time_unit) {
        return GetArrowType(primitive.physical_type(), *primitive.logical_type(),
                            primitive.type_length(), int96_arrow_time_unit);
    }
} // namespace parquet::arrow
